
from traits.api             import Range, Array, Instance, Enum, on_trait_change, Bool, Button
from traitsui.api           import View, Item, HGroup #, EnumEditor
from enable.api             import ComponentEditor
from chaco.api              import ArrayPlotData, Plot 
from traitsui.menu          import Action, Menu, MenuBar
from traitsui.file_dialog   import save_file

from datetime                           import datetime
from threading                          import currentThread

import numpy as np

from hardware.TimeTagger    import Pulsed, Countrate

from tools.emod             import FreeJob

from tools.utility          import GetSetItemsHandler, GetSetItemsMixin

class AutocorrelationHandler( GetSetItemsHandler ):
    
    def save_plot(self, info):
        filename = save_file(title='Save Plot')
        if filename is '':
            return
        else:
            if filename.find('.png')==-1:
                filename=filename+'.png'
            info.object.save_plot(filename)
    
class Autocorrelation( FreeJob, GetSetItemsMixin):
    
    WindowSize = Range(low=1, high=10000, value=500, desc='WindowSize in the time domain  [ns]', label='WindowSize')
    bin_width = Range(low=0.01, high=1000., value=1., desc='bin width  [ns]', label='bin width')
    chan1 = Enum(0,1,2,3,4,5,6,7, desc="the trigger channel", label="Channel 1")
    chan2 = Enum(0,1,2,3,4,5,6,7, desc="the signal channel", label="Channel 2")
    remove_zero = Bool(desc="remove the very high value at position zero", label="remove zero")
    restart = Button(desc="restart the measurement", label="restart")

    C = Array()
    T = Array()
    plot = Instance( Plot, transient=True )
    plot_data = Instance( ArrayPlotData, transient=True )
    
    get_set_items=['WindowSize', 'chan1', 'chan2', 'C', 'T', 'remove_zero']
    
    def __init__(self):
        super(Autocorrelation, self).__init__()
        self._update_pulsed()
        self.start()
        
    @on_trait_change('chan1,chan2,WindowSize,restart')    
    def _update_pulsed(self):
        self.p1 = Pulsed(int(np.round(self.WindowSize/self.bin_width)),int(np.round(self.bin_width*1000)),1,self.chan1,self.chan2)
        self.p2 = Pulsed(int(np.round(self.WindowSize/self.bin_width)),int(np.round(self.bin_width*1000)),1,self.chan2,self.chan1) 
        self.count_rate_1 = Countrate(self.chan1)
        self.count_rate_2 = Countrate(self.chan2)
        self.started = datetime.now()
    
    def _run(self):
        while True:
            data1 = self.p1.getData()
            data2 = self.p2.getData()
            if self.remove_zero:
                data = np.append(np.append(data1[0][-1:0:-1], 0),data2[0][1:])
            else:
                data = np.append(np.append(data1[0][-1:0:-1], max(data1[0][0],data2[0][0])),data2[0][1:])
        
            timediff = datetime.now() - self.started
            timediff = (timediff.seconds + timediff.days * 24 * 3600) 
            
            factor = self.count_rate_1.getData()*self.count_rate_2.getData()*1e-9*timediff
            if factor > 0:
                self.C = data / (self.count_rate_1.getData()*self.count_rate_2.getData()*1e-9*timediff)
            else:
                self.C = self._C_default()
        
            currentThread().stop_request.wait(1.0)
            if currentThread().stop_request.isSet():
                break

    def save_plot(self, filename):
        self.save_figure(self.plot, filename)
    
    def _remove_zero_default(self):
        return False
    
    def _C_default(self):
        return np.zeros((int(np.round(self.WindowSize/self.bin_width))*2-1,))
    
    def _T_default(self):
        return self.bin_width*np.arange(-int(np.round(self.WindowSize/self.bin_width))+1,int(np.round(self.WindowSize/self.bin_width)))
    
    def _T_changed(self):
        self.plot_data.set_data('t', self.T)

    def _C_changed(self):
        self.plot_data.set_data('y', self.C)
        self.plot.request_redraw()
    
    def _chan1_default(self):
        return 0
    
    def _chan2_default(self):
        return 1 
        
    def _WindowSize_changed(self):
        self.C = self._C_default()
        self.T = self._T_default()
        
    def _plot_data_default(self):
        return ArrayPlotData(t=self.T, y=self.C)
    
    def _plot_default(self):
        plot = Plot(self.plot_data, width=500, height=500, resizable='hv')
        plot.plot(('t','y'), type='line', color='blue')        
        return plot
    
    traits_view = View( 
        Item('plot', editor=ComponentEditor(), show_label=False),
        HGroup(
            Item('WindowSize'),
            Item('chan1'),
            Item('chan2'),
            Item('remove_zero'),
            Item('restart', show_label=False)),
        menubar = MenuBar(
            Menu(
                Action(action='save', name='Save (.pyd or .pys)'),
                Action(action='load', name='Load'),
                Action(action='save_plot', name='Save Plot (.png)'),
                Action(action='_on_close', name='Quit'),
            name='File')),
        title='Autocorrelation', width=900, height=800, buttons=[], resizable=True,
        handler=AutocorrelationHandler)

